import cv2
import numpy as np


# 图像骨架化处理
def skeletonize_image(img):
    skeleton = np.zeros(img.shape, np.uint8)
    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3,3))
    done = False
    while not done:
        eroded = cv2.erode(img, element)
        temp = cv2.dilate(eroded, element)
        temp = cv2.subtract(img, temp)
        skeleton = cv2.bitwise_or(skeleton, temp)
        img = eroded.copy()
        zeros = cv2.countNonZero(img)
        done = zeros == 0
    return skeleton

# 经pre_test文件多次尝试，设置图像预处理流程为：直方图自适应均衡化(图像增强)--->局部阈值法(二值化)
def preProcess(img):
    # 1.创建自适应直方图均衡化器
    # clipLimit: 这是一个阈值参数，用于控制对比度的限制。它指定了在直方图均衡化过程中对比度的最大增强量。
    # tileGridSize: 这是一个元组参数，用于指定图像被分成的小块的大小。
    # print(img.shape)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    # 2.应用自适应直方图均衡化，效果比较好
    adaptive_equalized_image = clahe.apply(img)
    # cv2.imshow('Adaptive Equalized Fingerprint', adaptive_equalized_image)
    # 局部阈值法,效果比较好
    thresholded_image = cv2.adaptiveThreshold(adaptive_equalized_image, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 2)
    out_image = cv2.adaptiveThreshold(thresholded_image, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 5, 2)
    # out_image = cv2.fastNlMeansDenoising(out_image, None, 5, 3, 7)
    # print(out_image.shape)
    # skeletionize = skeletonize_image(out_image)
    # return skeletionize
    return out_image


# 测试代码
if __name__ == '__main__':
    # 读取指纹图像
    fingerprint_img = cv2.imread('./data/test/1_1.tif', cv2.IMREAD_GRAYSCALE)
    result_img = preProcess(fingerprint_img)
    cv2.imshow('Original Fingerprint', fingerprint_img)
    cv2.imshow('PreProcess Fingerprint', result_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()